(*
  Copyright (c) 2009 Barry Schwartz

  Permission is hereby granted, free of charge, to any person
  obtaining a copy of this software and associated documentation
  files (the "Software"), to deal in the Software without
  restriction, including without limitation the rights to use,
  copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the
  Software is furnished to do so, subject to the following
  conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  OTHER DEALINGS IN THE SOFTWARE.
*)

(*-----------------------------------------------------------------------*)

(* A dynamic trie |t 'a| implements a map |list elt -> 'a| that supports adding new elements. *)

(*-----------------------------------------------------------------------*)

module type OrderedType = Map.OrderedType
  
module type S =
sig
  type elt
  type 'a t
    
  val empty : 'a t
  val is_empty : 'a t -> bool
  val prefix : 'a t -> elt -> 'a t
  val root_value : 'a t -> 'a option
  val depth : 'a t -> int
  val find_array : elt array -> 'a t -> 'a
  val mem_array : elt array -> 'a t -> bool
  val find_list : elt list -> 'a t -> 'a
  val mem_list : elt list -> 'a t -> bool
  val generic_lookup :
    (('a t -> elt -> 'a t) -> 'a t -> 'b -> 'a t) ->
    'b -> 'a t -> 'a option
  val generic_lookup_prefix :
    (((('a t) * ('a t)) -> elt -> (('a t) * ('a t))) ->
       (('a t) * ('a t)) -> 'b -> (('a t) * ('a t)))
    -> 'b -> 'a t -> 'a option
  val lookup_array : elt array -> 'a t -> 'a option
  val lookup_list : elt list -> 'a t -> 'a option
  val lookup_prefix_array : elt array -> 'a t -> 'a option
  val lookup_prefix_list : elt list -> 'a t -> 'a option
  val add_array : elt array -> 'a -> 'a t -> 'a t
  val remove_array : elt array -> 'a t -> 'a t
  val add_list : elt list -> 'a -> 'a t -> 'a t
  val remove_list : elt list -> 'a t -> 'a t
  val merge : 'a t -> 'a t -> 'a t
  val map : ('a -> 'b) -> 'a t -> 'b t
  val mapi : (elt array -> 'a -> 'b) -> 'a t -> 'b t
  val iter : (elt array -> 'a -> unit) -> 'a t -> unit
  val fold : (elt array -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
end                                     (* module type S *)

(*-----------------------------------------------------------------------*)
  
module Make (Ord : OrderedType) =
  struct
    module ChildMap = Map.Make(Ord)
      
    type elt = Ord.t
    
    type 'a t = { children : ('a t) ChildMap.t; data : 'a option }
    
    let empty = { children = ChildMap.empty; data = None; }
      
    let is_empty trie =
      (trie.data = None) && (ChildMap.is_empty trie.children)
      
    let prefix trie char =
      try ChildMap.find char trie.children
      with | Not_found -> empty
      
    let root_value trie = trie.data
      
    let generic_lookup fold str trie = root_value (fold prefix trie str)
      
    (*
      For efficiency reasons the following definitions are expanded:

      value lookup_array str trie = generic_lookup Array.fold_left str trie;
      value lookup_list  str trie = generic_lookup List.fold_left  str trie;
    *)
    let lookup_array str trie =
      let rec iter i trie =
        if i >= Array.length str then
          trie.data
        else
          try
            let next = ChildMap.find str.(i) trie.children in
              iter (i + 1) next
          with
            | Not_found -> None
      in
        iter 0 trie
      
    let rec lookup_list str trie =
      match str with
        | [] -> trie.data
        | c :: cs ->
            try
              let next = ChildMap.find c trie.children in
                lookup_list cs next
            with
              | Not_found -> None
      
    (*
      |lookup_prefix <str> <trie>| returns the data associated to the longest prefix of <str> that
      is found in <trie>.
    *)
    let generic_lookup_prefix fold str trie =
      root_value
        begin
          fst 
            begin
              fold (fun (last, t) c ->
                      let next = prefix t c in
                        match next.data with
                          | Some _ -> (next, next)
                          | None -> (last, next))
                (trie, trie)
                str
            end
        end

    (*
      For efficiency reasons the following definitions are expanded:

      value lookup_prefix_array str trie = generic_lookup_prefix Array.fold_left str trie;
      value lookup_prefix_list  str trie = generic_lookup_prefix List.fold_left  str trie;
    *)
    let lookup_prefix_array str trie =
      let rec iter i trie data =
        if i >= (Array.length str)
        then
          begin
            match trie.data with
              | Some _ -> trie.data
              | None -> data
          end
        else
          try
            let next = ChildMap.find str.(i) trie.children in
              match next.data with
                | Some _ -> iter (i + 1) next next.data
                | None -> iter (i + 1) next data
          with
            | Not_found -> data
      in
        iter 0 trie None
      
    let lookup_prefix_list str trie =
      let rec iter trie str data =
        match str with
          | [] ->
              begin
                match trie.data with
                  | Some _ -> trie.data
                  | None -> data
              end
          | c :: cs ->
              try
                let next = ChildMap.find c trie.children in
                  match next.data with
                    | Some _ -> iter next cs next.data
                    | None -> iter next cs data
              with
                | Not_found -> data
      in
        iter trie str None
      
    let find_array str trie =
      match lookup_array str trie with
        | Some x -> x
        | None -> raise Not_found
      
    let mem_array str trie =
      match lookup_array str trie with
        | Some _ -> true
        | None -> false
      
    let find_list str trie =
      match lookup_list str trie with
        | Some x -> x
        | None -> raise Not_found
      
    let mem_list str trie =
      match lookup_list str trie with
        | Some _ -> true
        | None -> false
      
    let add_array str data trie =
      let rec iter i trie =
        if i >= Array.length str then
          { trie with data = Some data }
        else
          let next =
            try ChildMap.find str.(i) trie.children
            with
              | Not_found -> empty
          in
            { trie with
                children = ChildMap.add str.(i) (iter (i + 1) next) trie.children }
      in
        iter 0 trie
      
    let remove_array str trie =
      let rec iter i trie =
        if i >= Array.length str then
          { trie with data = None }
        else
          try
            let next = ChildMap.find str.(i) trie.children in
              { trie with
                  children = ChildMap.add str.(i) (iter (i + 1) next) trie.children }
          with
            | Not_found -> trie
      in
        iter 0 trie
      
    let rec add_list str data trie =
      match str with
        | [] -> { trie with data = Some data }
        | c :: cs ->
            let next =
              try
                ChildMap.find c trie.children
              with
                | Not_found -> empty
            in
              { trie with
                  children = ChildMap.add c (add_list cs data next) trie.children }
      
    let rec remove_list str trie =
      match str with
        | [] -> { trie with data = None }
        | c :: cs ->
            (try
               let next = ChildMap.find c trie.children in
                 { trie with
                     children = ChildMap.add c (remove_list cs next) trie.children }
             with
               | Not_found -> trie)
      
    let merge t1 t2 =
      { data =
          begin
            match t2.data with
              | None -> t1.data
              | Some _ -> t2.data
          end;
        children =
          if ChildMap.is_empty t1.children then
            t2.children
          else
            ChildMap.fold (fun c x children -> ChildMap.add c x children)
              t2.children t1.children   }
      
    let rec map f trie =
      { children = ChildMap.map (map f) trie.children;
        data =
          match trie.data with
            | Some x -> Some (f x)
            | None -> None  }
      
    let mapi f trie =
      let rec iter str trie =
        { children = ChildMap.mapi (fun c t -> iter (c :: str) t) trie.children;
          data =
            match trie.data with
              | Some x -> Some (f (XList.rev_to_array str) x)
              | None -> None     }
      in
        iter [] trie
      
    let rec iter f trie =
      let rec loop str trie =
        begin
          match trie.data with
            | Some x -> f (XList.rev_to_array str) x
            | None -> ()
        end;
        ChildMap.iter (fun c t -> loop (c :: str) t) trie.children
      in
        loop [] trie
      
    let fold f trie e =
      let rec iter str trie e =
        let x =
          match trie.data with
            | Some y -> f (XList.rev_to_array str) y e
            | None -> e
        in
          ChildMap.fold (fun c t y -> iter (c :: str) t y) trie.children x
      in
        iter [] trie e
      
    let rec depth trie =
      if ChildMap.is_empty trie.children then
        0
      else
        let d = ChildMap.fold (fun _ t d -> max d (depth t)) trie.children 0 in
          d + 1
      
  end                                   (* module Make *)
  
(*-----------------------------------------------------------------------*)
